#include "LL1Parser.hpp"
#include <string>
#include <utility>
#include <vector>

std::unordered_set<std::string> LL1Parser::get_blank_set() {
  std::unordered_set<std::string> res;
  bool changed = true;
  while (changed) {
    auto original_size = res.size();
    for (const auto &item : defines) {
      if (item.rightItem.front() == "") {
        res.insert(item.leftItem);
        continue;
      }
      bool has_not_blank = false;
      for (const auto &symbol : item.rightItem) {
        if (!res.count(symbol)) {
          has_not_blank = true;
          break;
        }
      }
      if (!has_not_blank) {
        res.insert(item.leftItem);
      }
    }

    if (res.size() > original_size) {
      changed = true;
    } else {
      changed = false;
    }
  }

  blank_set = std::move(res);
  return blank_set;
}

void LL1Parser::mark_follow_idx() {
  auto blank_set = get_blank_set();
  for (auto &item : defines) {
    auto &right_item = item.rightItem;
    for (int i = right_item.size() - 1; i >= 0; --i) {
      if (terminals.count(right_item[i]) || !blank_set.count(right_item[i])) {
        item.follow_idx = i;
        break;
      }
    }
  }
}

std::unordered_map<std::string, std::unordered_set<std::string>>
LL1Parser::get_first_set() {
  bool changed = true;
  std::unordered_map<std::string, std::unordered_set<std::string>> res;
  for (const auto &item : defines) {
    res[item.leftItem] = std::unordered_set<std::string>{};
  }

  while (changed) {
    changed = false;
    for (const auto &item : defines) {
      auto front_item = item.rightItem.front();
      auto &target = res[item.leftItem];
      auto size = target.size();
      if (terminals.count(front_item)) {
        target.insert(front_item);
      } else {
        target.insert(res[front_item].begin(), res[front_item].end());
      }
      if (target.size() > size) {
        changed = true;
      }
    }
  }

  for (const auto &item : terminals) {
    res[item] = {item};
  }

  first_set = std::move(res);
  return first_set;
}

std::unordered_map<std::string, std::unordered_set<std::string>>
LL1Parser::get_follow_set() {
  get_first_set();
  mark_follow_idx();

  bool changed = true;
  std::unordered_map<std::string, std::unordered_set<std::string>> res;
  for (const auto &item : defines) {
    res[item.leftItem] = std::unordered_set<std::string>{};
  }

  while (changed) {
    std::unordered_map<std::string, size_t> size_map;
    for (const auto &kv : res) {
      size_map[kv.first] = kv.second.size();
    }

    //
    for (const auto &item : defines) {
      auto &right_item = item.rightItem;
      for (int i = 0; i < item.rightItem.size(); ++i) {
        auto &this_item = right_item[i];

        if (i >= item.follow_idx && !terminals.count(this_item)) {
          res[this_item].insert(res[item.leftItem].begin(),
                                res[item.leftItem].end());
        }
        if (i == 0) {
          continue;
        }
        auto &pre_item = right_item[i - 1];
        if (terminals.count(pre_item)) {
          continue;
        }

        if (terminals.count(this_item)) {
          res[pre_item].insert(this_item);
        } else {
          res[pre_item].insert(first_set[this_item].begin(),
                               first_set[this_item].end());
        }
      }

      res[item.leftItem].insert(END_SYMNOL);
      if (!terminals.count(right_item.back())) {
        res[right_item.back()].insert(END_SYMNOL);
      }
    }
    //

    changed = false;
    for (const auto &kv : res) {
      if (kv.second.size() > size_map[kv.first]) {
        changed = true;
        break;
      }
    }
  }

  for (auto &kv : res) {
    kv.second.erase("");
  }
  follow_set = std::move(res);
  return follow_set;
}

bool LL1Parser::can_be_blank(std::vector<std::string> right_item) {
  if (right_item.front() == "") {
    return true;
  }
  for (const auto &item : right_item) {
    if (!blank_set.count(item)) {
      return false;
    }
  }
  return true;
}

std::unordered_map<int, std::unordered_set<std::string>>
LL1Parser::get_select_set() {
  get_follow_set();

  std::unordered_map<int, std::unordered_set<std::string>> res;
  for (int i = 0; i < defines.size(); ++i) {
    const auto &tmp_item = defines[i].rightItem.front();
    if (can_be_blank(defines[i].rightItem)) {
      res[i] = follow_set[defines[i].leftItem];
      if (!terminals.count(tmp_item)) {
        auto to_insert = first_set[tmp_item];
        to_insert.erase("");
        res[i].insert(to_insert.begin(), to_insert.end());
      }
    } else {
      res[i] = first_set[tmp_item];
    }
  }

  select_set = std::move(res);
  return select_set;
}

std::unordered_map<std::string, std::vector<std::string>>
LL1Parser::get_table() {
  get_follow_set();

  std::unordered_map<std::string, std::vector<std::string>> res;
  for (const auto &kv : select_set) {
    auto &valid_set = kv.second;
    auto &define_item = defines[kv.first];
    for (const auto &item : valid_set) {
      res[define_item.leftItem + item] = define_item.rightItem;
    }
  }
  return res;
}
